#pragma once
#include "afxmt.h"

#define WDF_SECTION_NAME L"dnsfilter_Device.NT.Wdf"
#define SERVICE_SECTION_NAME _T("dnsfilter_Device.NT.Services")
#define DRIVER_SERVICE_NAME _T("dnsfilter")
#define DEVICE_KEY _T("SYSTEM\\CurrentControlSet\\Enum\\Root\\LEGACY_DNSFILTER")
#define DEVICE_KEY_INSTANCE _T("Root\\LEGACY_DNSFILTER\\0000")
#define DEVICE_ID_HARDWARE _T("Root\\dnsfilter")
#define SERVICE_DELETE_SECTION_NAME _T("dnsfilter_Device.NT.Services.Delete")

#define DNS_PORT    53

inline u_long hostAddress(
    u_long nOctet1,
    u_long nOctet2,
    u_long nOctet3,
    u_long nOctet4
    )
{
    u_long  nRet;
    nRet = (nOctet4&0xFF) + ((nOctet3&0xFF) << 8) + ((nOctet2&0xFF) << 16) + ((nOctet1&0xFF) << 24);
    return nRet;
}

inline u_long netAddress(
    u_long nOctet1,
    u_long nOctet2,
    u_long nOctet3,
    u_long nOctet4
    )
{
    u_long  nRet;
    nRet = (nOctet1&0xFF) + ((nOctet2&0xFF) << 8) + ((nOctet3&0xFF) << 16) + ((nOctet4&0xFF) << 24);
    return nRet;
}

typedef struct _DNS_APPROVAL_CONFIG
{
    BOOLEAN     m_bDiscard;
    BOOLEAN     m_bDataModify;
    UINT32      m_nDataSize;
} DNS_APPROVAL_CONFIG, *PDNS_APPROVAL_CONFIG;

// CPublicDNSPluginServiceApp

typedef enum ENStartMode
{
    enAutomatic,
    enManual
};

class CDnsResorceRecord;
class CMainWorkerThread;
class CPublicDNSPluginServiceApp : public CWinApp
{

public:
	CPublicDNSPluginServiceApp();

public:
	virtual BOOL InitInstance();
	virtual int ExitInstance();

protected:
	DECLARE_MESSAGE_MAP()
public:
    bool m_bInstallService;
    bool m_bUninstallService;
    bool m_bInstallDriver;
    bool m_bUninstallDriver;
    bool m_bRestart;
    bool m_bParseError;
    ENStartMode m_enStartMode;
    bool m_bUsageOut;
    int m_nRetCode;
    CString m_sModuleFileName;
    CString m_sModuleFolder;
    CString m_sInfFilePath;
    bool m_bRebootNeed;
    bool m_bTestModeInstalled;
    bool m_bRootCertExist;
    bool m_bRootCertInstalled;
    bool m_bTrustedPublisherCertExist;
    bool m_bTrustedPublisherCertInstalled;
    CString m_sCoInstallerPath;
    CString m_sDriversFolder;
    HMODULE m_hCoInatallerModule;
    PFN_WDFPREDEVICEINSTALL m_pfnWdfPreDeviceInstall;
    PFN_WDFPOSTDEVICEINSTALL m_pfnWdfPostDeviceInstall;
    PFN_WDFPREDEVICEINSTALLEX m_pfnWdfPreDeviceInstallEx;
    PFN_WDFPREDEVICEREMOVE m_pfnWdfPreDeviceRemove;
    PFN_WDFPOSTDEVICEREMOVE m_pfnWdfPostDeviceRemove;
    virtual int Run();
    static void WINAPI ServiceMain(
        DWORD dwArgc,
        LPTSTR* lpszArgv
        );
    static void ReportEventInvoke(
        WORD wType,
        DWORD dwEventID,
        WORD cInserts,
        LPCTSTR* szMsgArray
        );
    void mfServiceMain(
        DWORD dwArgc,
        PTSTR* lpszArgv
        );
    SERVICE_STATUS m_stServiceStatus;
    SERVICE_STATUS_HANDLE m_hServiceStatusHandle;
    static void WINAPI ServiceCtrlHandler(DWORD fdwControl);
    DWORD ServiceInitialization(
        DWORD argc,
        LPTSTR* argv,
        DWORD* specificError
        );
    void mfServiceCtrlHandler(DWORD fdwControl);
    CString m_sDnsZoneName;
    CDnsResorceRecord* m_pDnsZoneRecord;
    int m_nPublicIpCount;
    UINT32* m_pPublicIpArray;
    int m_nPublicRangeCount;
    FWP_RANGE0* m_pPublicRangeArray;
    int m_nPublicSubnetCount;
    FWP_V4_ADDR_AND_MASK* m_pPublicSubnetArray;
    int m_nPrivateIpCount;
    UINT32* m_pPrivateIpArray;
    int m_nPrivateRangeCount;
    FWP_RANGE0* m_pPrivateRangeArray;
    int m_nPrivateSubnetCount;
    FWP_V4_ADDR_AND_MASK* m_pPrivateSubnetArray;
    CMainWorkerThread* m_pWorkerThread;
    CMutex m_oWorkerThreadLock;
    CEvent m_oInitWorkerThreadDone;
    UINT32 m_dwCheckPoint;
    void RestartService(void);
    CString m_sDriverFolder;
    BYTE m_abAdminsSidBuffer[SECURITY_MAX_SID_SIZE];
    BYTE m_abWorldSidBuffer[SECURITY_MAX_SID_SIZE];
    ACL* m_pDeviceKeyDacl;
    SID* m_pAdminsSid;
    DWORD m_dwAdminsSidSize;
    SID* m_pWorldSid;
    DWORD m_dwWorldSidSize;
    HANDLE m_hProcessToken;
    bool SetProcessPrivilege(
        LPCTSTR lpszPrivilege,
        bool bEnablePrivilege
        );
    void mfDeleteKey(
        HKEY hRootKey,
        LPCTSTR szKeyPath
        );
	CString m_sOEMInfFile;
    int m_nPublicIPV6AddrCount;
    int m_nPrivateIPV6AddrCount;
    IPV6_ADDRESS_STRUCT* m_pPublicIPV6AddrArray;
    IPV6_ADDRESS_STRUCT* m_pPrivateIPV6AddrArray;
    void ShowDebugPromptDlg();
    CEvent m_oMainWorkerThreadDoneEvent;
};


extern CPublicDNSPluginServiceApp theApp;
